import time

import pytest
import yaml
from selenium import webdriver

class CookieLogin:
    def __init__(self):
        self.driver = webdriver.Chrome()
        self.driver.implicitly_wait(10)
        self.driver.get("https://work.weixin.qq.com/wework_admin/loginpage_wx")

    def set_cookie(self):
        time.sleep(10)
        cookies = self.driver.get_cookies()
        try :
            with open("cookie.yaml", "w") as f:
                yaml.safe_dump(cookies, f)
        except Exception as e:
            print(f"写入文件时出错: {e}")
    def get_cookie(self):
        try :
            with open("cookie.yaml", "r")as f :
                cookie = yaml.safe_load(f)
                for cookie in cookie:
                    self.driver.add_cookie(cookie)
        except Exception as e:
            print(f"读入文件时出错: {e}")
        self.driver.get("https://work.weixin.qq.com/wework_admin/loginpage_wx")

    def close_driver(self):
        self.driver.close()

if __name__ == "__main__":
    cookie = CookieLogin()
    # cookie.set_cookie()
    cookie.get_cookie()
    time.sleep(5)
    cookie.close_driver()